from environment import PendulumEnvironment
from controller.agent import Agent
import time

PENDULUM_STATE_DIM = 2
PENDULUM_ACTION_DIM = 2

env = PendulumEnvironment()
agent = Agent(state_dim=PENDULUM_STATE_DIM, action_dim=PENDULUM_ACTION_DIM,
              epsilon=1, fc1_dim=64, fc2_dim=64, max_size=10000, batch_size=32)
agent.load_models()

total_reward = 0
state = env.Reset()
for n in range(500):
    action = agent.choose_action((state[0], state[1]), isTrain=False)
    state = env.Step(action)
    total_reward += state[2]
bestReward = total_reward/500
print(f'Initial best reward:{bestReward}')

for episode in range(200):
    total_reward = 0
    done = False
    state = env.Reset()  # 返回角度和角速度2个参数
    for n in range(500):
        action = agent.choose_action(state, isTrain=True)  # 根据当前的2维状态选择1维动作
        observation = env.Step(action)  # 采取动作后，环境返回4维反馈，分别为：角度、角速度、奖励、是否结束
        agent.remember(state, action, observation[2], observation[0:2], observation[3])
        agent.learn()
        total_reward += observation[2]
        state = observation[0:2]
    currentWard = total_reward/500
    print(f'EP:{episode}, avg_reward:{currentWard} epsilon:{agent.epsilon}')
    if currentWard > bestReward:
        bestReward = currentWard
        agent.save_models()
    time.sleep(1)
